import re
import json
from typing import Dict, List, Any, Optional
from app.core.config import settings
from langchain_openai import ChatOpenAI
from app.core.logger import get_logger

logger = get_logger('agent.finance_ner')

class FinanceEntityRecognizer:
    """金融领域实体识别器，使用Qwen Finance模型进行专业实体识别"""
    
    def __init__(self):
        """初始化金融实体识别器"""
        # 初始化Qwen Finance模型
        self.finance_model = self._init_finance_model()
        
    def _init_finance_model(self):
        """初始化Qwen Finance模型"""
        try:
            dashscope_api_key = settings.DASHSCOPE_API_KEY
            dashscope_base_url = settings.DASHSCOPE_BASE_URL
            // 使用更新后的模型名称
            dashscope_finance_model = settings.DASHSCOPE_FINANCE_MODEL
            
            if dashscope_api_key:
                logger.info(f"Initializing Qwen Finance model: {dashscope_finance_model}")
                return ChatOpenAI(
                    model=dashscope_finance_model,
                    openai_api_key=dashscope_api_key,
                    openai_api_base=dashscope_base_url,
                    temperature=0.0,  # 实体识别任务使用较低温度确保一致性
                    max_tokens=1024
                )
            else:
                logger.warning("DASHSCOPE_API_KEY not found for Finance model")
                return None
        except Exception as e:
            logger.error(f"Failed to initialize Qwen Finance model: {e}")
            return None
    
    def extract_entities(self, query: str) -> Dict[str, Any]:
        """
        使用Qwen Finance模型提取金融实体
        
        Args:
            query: 用户查询
            
        Returns:
            包含识别实体的字典
        """
        if not self.finance_model:
            logger.warning("Qwen Finance model not available, falling back to basic entity extraction")
            return self._basic_entity_extraction(query)
        
        # 构建提示词
        prompt = self._build_ner_prompt(query)
        
        try:
            # 调用模型进行实体识别
            response = self.finance_model.invoke(prompt)
            
            # 解析响应
            entities = self._parse_ner_response(response)
            
            logger.info(f"Extracted entities: {entities}")
            return entities
        except Exception as e:
            logger.error(f"Failed to extract entities with Qwen Finance model: {e}")
            # 出错时回退到基础实体提取方法
            return self._basic_entity_extraction(query)
    
    def _build_ner_prompt(self, query: str) -> str:
        """
        构建实体识别提示词
        
        Args:
            query: 用户查询
            
        Returns:
            构建的提示词
        """
        prompt = f"""
你是一个专业的金融领域实体识别专家，请从以下用户查询中提取相关的金融实体。

用户查询: "{query}"

请严格按照以下JSON格式返回提取到的实体信息:
{{
    "company_names": ["提取到的公司名称列表"],
    "stock_codes": ["提取到的股票代码列表"],
    "fund_codes": ["提取到的基金代码列表"],
    "dates": ["提取到的日期列表，格式为YYYYMMDD"],
    "industries": ["提取到的行业名称列表"],
    "financial_terms": ["提取到的金融术语列表"],
    "metrics": ["提取到的财务指标列表，如：涨跌幅、市盈率、市净率等"],
    "amounts": ["提取到的金额或数量相关表述"],
    "confidence": "实体识别的置信度，0-1之间的数字"
}}

注意事项：
1. 如果某类实体没有识别到，请返回空列表
2. 公司名称请提取完整名称，如"贵州茅台股份有限公司"
3. 股票代码请包含市场标识，如"600519.SH"
4. 基金代码请包含完整代码，如"000001"
5. 日期请统一为YYYYMMDD格式
6. 置信度请根据识别确定程度给出，0表示完全不确定，1表示完全确定

只返回JSON格式的结果，不要添加其他说明。
"""
        return prompt
    
    def _parse_ner_response(self, response) -> Dict[str, Any]:
        """
        解析实体识别响应
        
        Args:
            response: 模型响应
            
        Returns:
            解析后的实体字典
        """
        try:
            # 确保响应是字符串
            if hasattr(response, 'content'):
                response_text = response.content
            else:
                response_text = str(response)
            
            # 清理响应文本，确保是有效的JSON
            response_text = response_text.strip()
            if response_text.startswith("```json"):
                response_text = response_text[7:]
            if response_text.endswith("```"):
                response_text = response_text[:-3]
            
            # 解析JSON
            entities = json.loads(response_text)
            return entities
        except Exception as e:
            logger.error(f"Failed to parse NER response: {e}")
            # 返回默认实体结构
            return {
                "company_names": [],
                "stock_codes": [],
                "fund_codes": [],
                "dates": [],
                "industries": [],
                "financial_terms": [],
                "metrics": [],
                "amounts": [],
                "confidence": 0.0
            }
    
    def _basic_entity_extraction(self, query: str) -> Dict[str, Any]:
        """
        基础实体提取方法（正则表达式）
        
        Args:
            query: 用户查询
            
        Returns:
            提取的实体字典
        """
        entities = {
            "company_names": [],
            "stock_codes": [],
            "fund_codes": [],
            "dates": [],
            "industries": [],
            "financial_terms": [],
            "metrics": [],
            "amounts": [],
            "confidence": 0.5  # 基础方法置信度较低
        }
        
        # 提取公司名称（简单模式）
        company_patterns = [
            r'([^，。,;；\s]+?[股份有限]{2,4}公司)',
            r'([^，。,;；\s]+?有限责任公司)'
        ]
        
        for pattern in company_patterns:
            matches = re.findall(pattern, query)
            if matches:
                entities["company_names"].extend(matches)
        
        # 提取股票代码
        stock_patterns = [
            r'([0-9]{6}\.[Ss][Hh])',  # 上交所股票代码格式，如 600000.SH
            r'([0-9]{6}\.[Ss][Zz])',  # 深交所股票代码格式，如 000001.SZ
            r'([0-9]{6})'  # 简单6位数字代码
        ]
        
        for pattern in stock_patterns:
            matches = re.findall(pattern, query)
            entities["stock_codes"].extend(matches)
        
        # 提取基金代码
        fund_pattern = r'[0-9]{6}'
        fund_matches = re.findall(fund_pattern, query)
        entities["fund_codes"].extend(fund_matches)
        
        # 提取日期
        date_patterns = [
            r'(\d{4}年\d{1,2}月\d{1,2}日)',
            r'(\d{4}-\d{1,2}-\d{1,2})',
            r'(\d{4}/\d{1,2}/\d{1,2})',
            r'(\d{8})'
        ]
        
        for pattern in date_patterns:
            matches = re.findall(pattern, query)
            entities["dates"].extend(matches)
        
        # 提取财务指标
        metrics = ["涨跌幅", "市盈率", "市净率", "ROE", "营收", "净利润", "毛利率", "净利率", "成交量", "成交额"]
        for metric in metrics:
            if metric in query:
                entities["metrics"].append(metric)
        
        return entities

# 全局实例
entity_recognizer = FinanceEntityRecognizer()